﻿using FlyIL;
using FlyIL.Emu;
using FlyIL.Lifter;
using FlyIL.Misc;
using FlyIL.Advance;
using FlyRemote.Base;
using FuckSp;
using System;
using System.Collections.Generic;
using System.IO;
using System.Linq;
using TestEmu;
using static FlyIL.OperandMaker;
using static FlyIL.AssemblerRegisters;

namespace TestFw
{
    internal class CheckCollect
    {
        IReader Reader;
        IReport Report;

        bool Is64;

        ReadGroup ReadGroup { get; set; }

        /// <summary>
        /// 所有的 handler ,序列化保存 
        /// </summary>
        public Dictionary<ulong, Context> handlers = new Dictionary<ulong, Context>();

        public Context Vcontext { get; private set; }


        public CheckCollect(bool is64, ReadGroup group, IReader reader)
        {
            this.Is64 = is64; Vcontext = new Context(is64); ReadGroup = group; this.Reader = reader;
        }




        ulong FindReadNext(Context handler)
        {
            /// 042A354F 这一处 handler 是什么 ???

            if (handler.EntryPoint == 0x0425BDAE)
                ;

            for (var node = handler.Last; node != null; node = handler.GetPreviousIfExist(node))
                /// 从 [esi+???]  或者从 esi 转移后的 寄存器  取常量  
                if (node.Value.OpLast == _dword_[R6D])
                {
                    return node.Value.Vip;
                }

            return 0;
        }


        /// <summary>
        /// 检测 handler 的类型 
        /// </summary>
        /// <param name="handler"></param>
        /// <param name="node"></param>
        /// <returns></returns>
        internal static HandleType CheckHandle(Context handler, out LinkedListNode<Inst> node)
        {
            node = null;


            bool read_vip = false;

            bool push_esi = false;

            for (var u = handler.Entry; u != null; u = handler.GetNextIfExist(u))
            {
                if (u.Value.Vip == 0x042888F2)
                    ;

                if (!push_esi)
                    push_esi = u.Value.IsMatch(Opt.push, _regs_[R6D], null);

                if (!push_esi)
                {

                    /// 如果是 mov esi, xxxx 或者 pop esi ,则是 跳转  
                    var a = u.Value.IsMatch(Opt.mov, _regs_[AssemblerRegisters.R6D], null);
                    var b = u.Value.IsMatch(Opt.pop, _regs_[AssemblerRegisters.R6D], null);

                    if (!read_vip && (a || b))
                        return HandleType.jmp;
                }


                {

                    /// 如果存在 push [esi] 或者 mov xxx , [esi] ,则可能是计算下一处 handler 
                    var a = u.Value.IsMatch(Opt.push, _dword_[AssemblerRegisters.R6D], null);
                    var b = u.Value.IsMatch(Opt.mov, _regs_[AssemblerRegisters.Any], _dword_[AssemblerRegisters.R6D]);
                    if (!read_vip && a || b)
                        read_vip = true;

                }

                if (!read_vip && u.Value.IsJxx)
                {
                    node = u;

                    /// 如果 有rep 指令, 则是 checkvsp , 否则是 jcc

                    return handler.Codes.Any(o => o.HavePrefix) ? HandleType.checkVsp : HandleType.jcc;
                }
            }

            return HandleType.normal;
        }



        Context prepareHandler(bool is64, ulong addr, IReader reader)
        {
            /// todo: 首个片段 使用 checkRet , 会造成栈错位 ，禁用后无此问题
            var handler = new Context(is64, addr);

            /// 屏蔽这些 计算 
            handler.OptimizerOption.IgnoreOpts.Add(Opt.adc);
            var lifter = new Lifters.Intel(is64, reader);
            var config = new Lifters.IntelOption { CheckRet = false, CallAsjmp = true, TrustJcc = true };

            lifter.Collect(handler, addr, config);

            ulong read_next = FindReadNext(handler);

            LinkedListNode<Inst> read_node = null;
            if (read_next != 0)
            {

                handler.Marks["next"] = read_next;
                read_node = handler[read_next];
            }

            /// 如果是 vjcc 需要把 未跳转的代码 也收集过来
            if (CheckHandle(handler, out LinkedListNode<Inst> jno) >= HandleType.jcc)
            {
                lifter.CollectNext(handler, jno.Value.Vip, config);
            }

            /**************************
             * 在收集后 即对 handler 化简 ，即可验证化简的正确性 ，又可提升效率 
             * ***********************/
            if (handler.EntryPoint == 0x000000000426bdb3)
                ;

            if (read_node != null)
                handler.SimplifyFlow(handler.Entry, read_node);

            ///// 移除这些操作数 不会影响 cpu 的执行
            List<Operand> remove2 = new() { _regs_[R1D], _regs_[CF], _regs_[R3D], _regs_[EFlag] };

            remove2.Add(_regs_[R4D]);

            /// 不能移除    edi  虚拟栈  esi 虚拟码地址 


            handler.RemoveDeadBatch(handler.GetPreviousIfExist(read_node), remove2);
            //handler.DumpFlow();


            handlers[addr] = handler;

            return handler;
        }




        /// <summary>
        /// 
        /// </summary>
        /// <param name="is64"></param>
        /// <param name="reader"></param>
        /// <param name="report"></param>
        /// <param name="hAddr"> handler 地址 </param>
        /// <param name="vip">入口的esi值，即虚拟地址</param>
        public void Collect(ulong hAddr, ulong vip = 0)
        {

            // 收集到总的上下文中  
            var con = Vcontext;

            Cpu cpu = new(Is64, Reader);


            var ReportPath = "C:\\Users\\lyo\\Desktop\\aaa.txt";
            File.Delete(ReportPath);


            cpu.SetReg64(RegCode.r6, vip);

            ulong handle_addr = hAddr;
            ulong cip = 0;

            bool removeMore = false;

            for (int i = 0; ; handle_addr = cpu.NextIp, i++)
            {
                cip = cpu.GetReg64(RegCode.r6);


                Context handler;
                char firstFound = '+';
                string text = null;

                if (i == 789)
                    ;

                if (!handlers.ContainsKey(handle_addr))
                {
                    handler = prepareHandler(Is64, handle_addr, Reader);
                }
                else
                {
                    firstFound = ' ';
                    handler = handlers[handle_addr];
                }
                handler.DumpFlow();


                var h_type = CheckHandle(handler, out LinkedListNode<Inst> jnode);


                // 当前这一串代码的 vip 
                text = string.Format("第{0:d4}段, {1} VIP= {2:X8} , HANDLE= {3:X8} , {4}\r\n", i, firstFound, cip, handle_addr, h_type.ToString());
                Console.WriteLine(text);


                File.AppendAllText(ReportPath, text);


                /// 必须随时 释放资源，否则后面无法继续创建 
                using (UnicornWrapper wrapper = new UnicornWrapper(Is64, ReadGroup))
                {

                    wrapper.CopyFrom(cpu);

                    cpu.Run(handler, handler.Entry, null, null);

                    /// 如果是 vinit ,那么修正 cip 地址为 下一个  cip -4  
                    if (cip == 0)
                        cip = cpu.GetReg64(RegCode.r6) - 4;


                    /// 如果使用 ClearNop ， 会使地址失准 
                    wrapper.Run(handle_addr, cpu.NextIp, 0);


                    if (!wrapper.CompareStack(cpu))
                    {
                        if (i != 0x50 /* 此处 unicorn 存在问题， 因为 bsf 已经将 所有flags 清理， 而unicorn 的标志有问题 */

                            )
                        {

                            /// todo: 首个片段 使用 checkRet , 会造成栈错位 ，禁用后无此问题
                            var orign = new Context(false, handler.EntryPoint);

                            var lifter = new Lifters.Intel(false, Reader);
                            var config = new Lifters.IntelOption { CheckRet = false, CallAsjmp = true, TrustJcc = true };

                            lifter.Collect(orign, handler.EntryPoint, config);


                            orign.DumpFlow();
                            handler.DumpPath = "d:\\dump-il2.txt";
                            handler.DumpFlow();



                            throw new Exception();
                        }
                    }
                }




                if (i == 32)
                    ;

            }
        }


    }
}
